import streamlit as st
import tempfile
import os
import json
from datetime import datetime

# 页面配置
st.set_page_config(
    page_title="知识库管理演示版", 
    page_icon="📚", 
    layout="wide"
)

st.title("📚 知识库管理及检索系统 (演示版)")

# 模拟的知识库数据
if "knowledge_base" not in st.session_state:
    st.session_state.knowledge_base = [
        {
            "id": "doc1",
            "title": "LLM自主代理简介",
            "content": "以大语言模型（LLM）为核心控制器构建代理是一个很酷的概念。几个概念验证演示，如AutoGPT、GPT-Engineer和BabyAGI，都是鼓舞人心的例子。LLM的潜力超越了生成优秀的文案、故事、论文和程序；它可以被框架化为一个强大的通用问题解决器。",
            "source": "网页",
            "created_at": "2024-01-15 10:30:00"
        },
        {
            "id": "doc2", 
            "title": "任务分解技术",
            "content": "任务分解是将复杂任务分解为更小、更易管理的步骤的技术。它可以通过思维链（CoT）提示等方法实现，指导模型'逐步思考'，利用更多测试时间计算将困难任务分解为更小更简单的步骤。",
            "source": "文档",
            "created_at": "2024-01-15 11:00:00"
        },
        {
            "id": "doc3",
            "title": "自我反思机制", 
            "content": "自我反思是允许自主代理通过改进过去的行动决策和纠正以前的错误来迭代改进的重要方面。它在现实世界任务中发挥关键作用，在这些任务中试错是不可避免的。",
            "source": "研究论文",
            "created_at": "2024-01-15 11:30:00"
        }
    ]

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# 侧边栏配置
with st.sidebar:
    st.header("🔧 系统配置")
    
    # 模拟API配置
    st.subheader("API 配置")
    api_configured = st.checkbox("模拟 API 已配置", value=True)
    
    # 知识库统计
    st.subheader("📊 知识库统计")
    st.metric("文档总数", len(st.session_state.knowledge_base))
    st.metric("聊天记录", len(st.session_state.chat_history))
    
    # 清空功能
    if st.button("🗑️ 清空知识库"):
        st.session_state.knowledge_base = []
        st.rerun()
    
    if st.button("🗑️ 清空聊天记录"):
        st.session_state.chat_history = []
        st.rerun()

# 主界面
col1, col2 = st.columns([1, 1])

with col1:
    st.header("📄 文档管理")
    
    # 文档上传模拟
    st.subheader("添加文档")
    
    with st.form("add_document"):
        doc_title = st.text_input("文档标题")
        doc_content = st.text_area("文档内容", height=150)
        doc_source = st.selectbox("文档来源", ["网页", "文档", "研究论文", "其他"])
        
        if st.form_submit_button("📥 添加到知识库"):
            if doc_title and doc_content:
                new_doc = {
                    "id": f"doc{len(st.session_state.knowledge_base) + 1}",
                    "title": doc_title,
                    "content": doc_content,
                    "source": doc_source,
                    "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                }
                st.session_state.knowledge_base.append(new_doc)
                st.success("✅ 文档已添加到知识库")
                st.rerun()
            else:
                st.error("请填写标题和内容")
    
    # 文件上传模拟
    st.subheader("上传文件")
    uploaded_files = st.file_uploader(
        "选择文件",
        type=['pdf', 'txt', 'md'],
        accept_multiple_files=True,
        key="uploaded_files"
    )
    
    if uploaded_files:
        for uploaded_file in uploaded_files:
            if st.button(f"📄 处理 {uploaded_file.name}"):
                # 模拟文件处理
                file_content = f"这是从文件 {uploaded_file.name} 中提取的内容示例。实际应用中这里会是真实的文件内容。"
                new_doc = {
                    "id": f"doc{len(st.session_state.knowledge_base) + 1}",
                    "title": uploaded_file.name,
                    "content": file_content,
                    "source": "上传文件",
                    "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                }
                st.session_state.knowledge_base.append(new_doc)
                st.success(f"✅ 文件 {uploaded_file.name} 已处理并添加到知识库")
                st.rerun()
    
    # 知识库内容展示
    st.subheader("📚 知识库内容")
    if st.session_state.knowledge_base:
        for i, doc in enumerate(st.session_state.knowledge_base):
            with st.expander(f"{doc['title']} ({doc['source']})"):
                st.text(f"ID: {doc['id']}")
                st.text(f"创建时间: {doc['created_at']}")
                st.text_area(f"内容", doc['content'], key=f"content_{i}", disabled=True)
                if st.button(f"🗑️ 删除", key=f"delete_{i}"):
                    st.session_state.knowledge_base.pop(i)
                    st.rerun()
    else:
        st.info("知识库为空，请添加文档")

with col2:
    st.header("🔍 知识检索")
    
    if not api_configured:
        st.warning("请先在侧边栏配置API")
    else:
        # 模拟的智能问答
        st.subheader("💬 智能问答")
        
        # 显示聊天历史
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
                if message["role"] == "assistant" and "sources" in message:
                    with st.expander("📚 参考文档"):
                        for source in message["sources"]:
                            st.markdown(f"**{source['title']}**")
                            st.text(source['content'][:200] + "...")
        
        # 用户输入
        if question := st.chat_input("请输入您的问题..."):
            # 显示用户消息
            with st.chat_message("user"):
                st.markdown(question)
            st.session_state.chat_history.append({"role": "user", "content": question})
            
            # 模拟检索和回答
            with st.chat_message("assistant"):
                with st.spinner("正在思考..."):
                    # 简单的关键词匹配检索
                    relevant_docs = []
                    question_lower = question.lower()
                    for doc in st.session_state.knowledge_base:
                        if any(keyword in doc['content'].lower() for keyword in question_lower.split()):
                            relevant_docs.append(doc)
                    
                    # 生成模拟回答
                    if relevant_docs:
                        answer = f"根据知识库中的相关文档，我找到了 {len(relevant_docs)} 个相关资料。"
                        
                        if "任务分解" in question or "分解" in question:
                            answer += "任务分解是一种重要的技术，它将复杂任务分解为更小、更易管理的步骤。通过思维链（CoT）提示等方法，可以指导模型逐步思考，有效解决困难问题。"
                        elif "自我反思" in question or "反思" in question:
                            answer += "自我反思是自主代理的关键能力，它允许系统通过改进过去的决策和纠正错误来持续改进。这在现实世界任务中特别重要。"
                        elif "LLM" in question or "语言模型" in question:
                            answer += "大语言模型（LLM）作为核心控制器可以构建强大的代理系统。它不仅能生成文本，还可以作为通用问题解决器使用。"
                        else:
                            answer += "这是一个很好的问题。基于当前知识库的内容，我可以为您提供相关信息。"
                        
                        st.markdown(answer)
                        
                        # 显示相关文档
                        with st.expander("📚 相关文档片段"):
                            for doc in relevant_docs[:3]:  # 最多显示3个相关文档
                                st.markdown(f"**{doc['title']}:**")
                                st.text(doc['content'][:300] + "...")
                                st.markdown("---")
                        
                        st.session_state.chat_history.append({
                            "role": "assistant", 
                            "content": answer,
                            "sources": relevant_docs[:3]
                        })
                    else:
                        answer = "抱歉，我在当前知识库中没有找到与您问题相关的信息。请尝试添加更多相关文档到知识库中。"
                        st.markdown(answer)
                        st.session_state.chat_history.append({"role": "assistant", "content": answer})

# 底部信息
st.markdown("---")
col1, col2 = st.columns([2, 1])

with col1:
    st.markdown("""
    ### 📝 使用说明（演示版）

    1. **添加文档**: 在左侧手动输入文档内容或模拟上传文件
    2. **智能问答**: 在右侧输入问题，系统会基于知识库内容回答
    3. **检索机制**: 使用简单的关键词匹配进行文档检索
    4. **管理功能**: 可以查看、删除知识库中的文档

    ### 🎯 演示特性
    - 模拟文档处理和存储
    - 基于关键词的简单检索
    - 交互式问答界面
    - 知识库管理功能
    """)

with col2:
    st.markdown("""
    ### 🔧 完整版功能
    - DeepSeek LLM 集成
    - Ollama 嵌入模型
    - PostgreSQL 向量存储
    - LangChain 检索链
    - 智能文档分割
    - 语义搜索
    """)

# 导出功能
st.markdown("### 💾 数据管理")
col1, col2 = st.columns(2)

with col1:
    if st.button("📤 导出知识库"):
        export_data = {
            "knowledge_base": st.session_state.knowledge_base,
            "export_time": datetime.now().isoformat()
        }
        st.download_button(
            label="下载 JSON 文件",
            data=json.dumps(export_data, ensure_ascii=False, indent=2),
            file_name=f"knowledge_base_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
            mime="application/json"
        )

with col2:
    uploaded_json = st.file_uploader("📥 导入知识库", type=['json'])
    if uploaded_json:
        try:
            import_data = json.loads(uploaded_json.read())
            if "knowledge_base" in import_data:
                st.session_state.knowledge_base.extend(import_data["knowledge_base"])
                st.success("✅ 知识库导入成功")
                st.rerun()
        except Exception as e:
            st.error(f"导入失败: {str(e)}")